import akshare as ak
import pandas as pd
import os
from datetime import datetime

class FinancialDataFetcher:
    """
    财务数据获取类
    负责从AKShare获取股票财务数据并进行初步处理
    """
    def __init__(self, data_dir="../data/raw"):
        self.data_dir = data_dir
        self.data_cache = {}
        # 确保数据存储目录存在
        os.makedirs(self.data_dir, exist_ok=True)
    
    def fetch_balance_sheet(self, stock_code, report_type='年报'):
        """
        获取资产负债表数据
        
        参数:
            stock_code (str): 股票代码，如'600000'或'sh600000'
            report_type (str): 报告类型，默认为'年报'
            
        返回:
            pandas.DataFrame: 处理后的资产负债表数据
        """
        print(f"开始获取股票 {stock_code} 的资产负债表数据...")
        
        # 确保股票代码格式正确（添加市场前缀）
        formatted_code = stock_code
        if not (stock_code.startswith('sh') or stock_code.startswith('sz')):
            # 根据股票代码判断市场
            if stock_code.startswith('6'):
                formatted_code = f'sh{stock_code}'
            else:
                formatted_code = f'sz{stock_code}'
        
        print(f"使用格式化后的股票代码: {formatted_code}")
        
        try:
            # 使用AKShare获取资产负债表数据
            print(f"调用AKShare API获取数据...")
            data = ak.stock_financial_report_sina(stock=formatted_code, symbol="资产负债表")
            
            if data is None or data.empty:
                print(f"AKShare返回空数据，请检查股票代码是否有效")
                return None
                
            print(f"成功获取原始数据! 数据形状: {data.shape}")
            
            # 数据初步处理
            processed_data = self._process_balance_sheet(data)
            
            if processed_data is None or processed_data.empty:
                print("数据处理后为空，请检查数据格式")
                return None
                
            print(f"数据处理完成! 处理后数据形状: {processed_data.shape}")
            
            # 缓存数据
            cache_key = f"{formatted_code}_{report_type}_balance_sheet"
            self.data_cache[cache_key] = processed_data
            
            return processed_data
        
        except Exception as e:
            import traceback
            print(f"获取资产负债表数据失败: {e}")
            print("详细错误信息:")
            traceback.print_exc()
            return None
    
    def _process_balance_sheet(self, data):
        """
        处理原始资产负债表数据
        
        参数:
            data (pandas.DataFrame): 原始资产负债表数据
            
        返回:
            pandas.DataFrame: 处理后的资产负债表数据
        """
        if data is None or data.empty:
            print("输入数据为空，无法处理")
            return None
        
        print(f"开始处理原始数据，列名: {list(data.columns)}")
        
        # 检查数据结构
        # 如果第一列是科目名称，需要设置为索引
        if len(data.columns) > 0:
            if '科目' in data.columns:
                print("检测到'科目'列，设置为索引")
                data.set_index('科目', inplace=True)
            elif data.columns[0] == '报告日期' or data.columns[0] == '报告日':
                print(f"检测到第一列为'{data.columns[0]}'，数据可能已经转置")
                # 数据可能已经转置，不需要额外处理
                pass
            else:
                print(f"未检测到标准列名，第一列为: {data.columns[0]}")
        
        # 确保数据格式正确
        # 将所有数值列转换为数值类型
        try:
            numeric_cols = []
            for col in data.columns:
                if col not in ['报表日期', '报告日期', '报告日', '科目', '类型', '更新日期', '更新日']:
                    print(f"尝试将列 '{col}' 转换为数值类型")
                    data[col] = pd.to_numeric(data[col], errors='coerce')
                    numeric_cols.append(col)
            
            if not numeric_cols:
                print("警告: 未找到可转换为数值的列")
                return data
                
            print(f"成功转换 {len(numeric_cols)} 列为数值类型")
            
            # 单位转换为亿元（假设原始数据单位为万元）
            for col in numeric_cols:
                data[col] = data[col] / 10000  # 万元转换为亿元
            
            print("单位已转换为亿元")
            
        except Exception as e:
            import traceback
            print(f"处理数据时出错: {e}")
            print("详细错误信息:")
            traceback.print_exc()
            # 如果处理过程中出错，返回原始数据，避免数据丢失
            return data
        
        return data
    
    def save_to_local(self, data, stock_code, report_type='年报'):
        """
        保存数据到本地CSV文件
        
        参数:
            data (pandas.DataFrame): 要保存的数据
            stock_code (str): 股票代码
            report_type (str): 报告类型
        """
        if data is None or data.empty:
            print("没有数据可保存")
            return
        
        # 创建文件名
        filename = f"{stock_code}_{report_type}_balance_sheet_{datetime.now().strftime('%Y%m%d')}.csv"
        file_path = os.path.join(self.data_dir, filename)
        
        # 保存到CSV
        data.to_csv(file_path, encoding='utf-8-sig')
        print(f"数据已保存到: {file_path}")
        
        return file_path

# 测试代码
if __name__ == "__main__":
    fetcher = FinancialDataFetcher()
    # 获取平安银行(000001)的资产负债表
    data = fetcher.fetch_balance_sheet("000001")
    if data is not None:
        print(data.head())
        # 保存到本地
        fetcher.save_to_local(data, "000001")